import sys
from types import FunctionType
import os, os.path, shutil
from os.path import join as opjoin
import base64
from tempfile import mkdtemp
import cPickle as pickle
from util import *

## Task template
taskrunner = \
"""
from types import FunctionType
import cPickle as pickle
import base64, os

if '_anycloud_payload' in task:
    file('__payload__.tar', 'wb').write(base64.standard_b64decode(task['payload']))
    os.system("tar -xf __payload__.tar")

ns = task['push']
for name,val in ns.iteritems():
    if val.__class__.__name__ == 'CannedFunction':
        ns[name] = FunctionType(val['code'], ns)

exec task['expression'] in ns
output = {
    'result': dict((k,ns.get(k)) for k in task['pull']),
    'error': None
}
pickle.dump(output, file('__result__.pickle', 'w'))
"""

class CannedFunction(dict):
    pass

## Task status
class TaskStatus:
    def __init__(self, name):
        self.name = name
    def __repr__(self):
        return "<TaskStatus %s>" % self.name

TASK_CREATED = TaskStatus('CREATED')
TASK_SUBMITTED = TaskStatus('SUBMITTED')
TASK_SUCCESS = TaskStatus('SUCCESS')
TASK_ERROR = TaskStatus('ERROR')

class TaskResult(object):
    def __init__(self, task, result, error):
        self.task = task
        self.result = result
        self.error = error

class StringTask(object):
    _last_id = 0

    def __init__(self, expression, pull=None, push=None, files=None, modules=None): 
        self.id = Task._new_id()
        self.expression = expression
        self.pull = pull if pull else ()
        self.push = push if push else {}
        self.files = files if files else []
        self.modules = modules if modules else []
        self.engineinfo = {}
        self.status = TASK_CREATED

        for name,val in self.push.iteritems():
            if isinstance(val, FunctionType):
                self.push[name].CannedFunction(code=val.func_code)
    
    @classmethod
    def _new_id(cls):
        cls._last_id += 1
        return cls._last_id

    def _export(self, outdir=None, singlefile=False):
        if not outdir:
            outdir = mkdtemp("anycloud")
        
        # export task params
        taskdict = dict(
            expression=self.expression,
            push=self.push,
            pull=self.pull,
        )
        
        if singlefile and (self.files or self.modules):
            payload = self._export_payload(outdir, as_string=True)
            taskdict.update({'__payload__':payload})
        else:
            self._export_payload(outdir)

        picklestr = pickle.dumps(taskdict)
        prefixstr = "import cPickle as pickle\ntask = pickle.loads(\"\"\"%s\"\"\")\n" % picklestr
        with file(opjoin(outdir, '__task__.py'), 'w') as tfile:
            tfile.write(prefixstr + taskrunner)

        self.wd = outdir
        return outdir

    def _export_payload(self, outdir=None, as_string=False):
        if not outdir:
            outdir = mkdtemp("anycloud")
        
        # export task modules and files
        for f in self.files:
            export_path(f, outdir)
        for m in self.modules:
            export_module(m, outdir)

        if as_string:
            # first tar it up, then read tarfile as binary file 
            oldcwd = os.getcwd()
            os.chdir(outdir)
            os.system("tar -cf __payload__.tar *")
            payload = file(opjoin(outdir, '__payload__.tar'), 'rb').read()
            payload = base64.standard_b64encode(payload)
            os.chdir(oldcwd)
            return payload
        
    def _process_result(self, wd):
        output = pickle.load(file(opjoin(wd, '__result__.pickle')))
        self.result = TaskResult(self, output['result'], output['error'])
        self.status = TASK_ERROR if output['error'] else TASK_SUCCESS

    def _cleanup(self):
        for f in os.listdir(self.wd):
            os.remove(os.path.join(self.wd, f))
        os.rmdir(self.wd)

class Task(StringTask):
    def __init__(self, func, args=None, kwargs=None, files=None, modules=None): 
        self.id = Task._new_id()
        self.func = func
        self.args = args or ()
        self.kwargs = kwargs or {}
        self.expression = "_anycloud_result = _anycloud_func(*_anycloud_args,**_anycloud_kwargs)"
        self.push = dict(
            _anycloud_func=CannedFunction(code=self.func.func_code),
            _anycloud_args=self.args,
            _anycloud_kwargs=self.kwargs
        )
        self.pull = ['_anycloud_result']
        self.files = files or []
        self.modules = modules or []
        self.engineinfo = {}
        self.status = TASK_CREATED
    
    def _process_result(self, wd):
        output = pickle.load(file(opjoin(wd, '__result__.pickle')))
        self.result = output['result']['_anycloud_result'] # TODO: error handling
        self.status = TASK_ERROR if output['error'] else TASK_SUCCESS

def map(func, seq, files=None, modules=None):
    return [Task(func, (s,), files=files, modules=modules) for s in seq]

